import torch
from tqdm import tqdm
from model import ReviewAnalyzeModel
from dataset import get_dataloader
from predict import predict_batch
from config import MODELS_DIR


def evaluate(model, dataloader, device):
    """ 评估 """
    model.eval()
    total_count = 0
    correct_count = 0

    for batch in tqdm(dataloader, desc="evaluate"):
        input_ids = batch['input_ids'].to(device)  # [batch_size, seq_len]
        attention_mask = batch['attention_mask'].to(device)  # [batch_size, seq_len]
        targets = batch['label'].to(device)  # [batch_size]

        probs = predict_batch(model, input_ids, attention_mask)  # [batch_size]

        for prob, target in zip(probs, targets):
            prob_label = 1 if prob > 0.5 else 0
            if prob_label == target:
                correct_count += 1
            total_count += 1

    accuracy = correct_count / total_count
    return accuracy


def run_evaluate():
    """ 评估流程 """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 设备

    # 创建并加载模型
    model = ReviewAnalyzeModel(freeze_bert=True).to(device)
    model.load_state_dict(torch.load(MODELS_DIR / 'model.pt'))

    # 测试数据
    dataloader = get_dataloader(train=False)

    # 执行评估
    accuracy = evaluate(model, dataloader, device)
    print(f'准确率：{accuracy:.4f}')


if __name__ == '__main__':
    run_evaluate()
